package com.li.springbootproject.config.requestlimit.bucket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author LiDong
 * @version 1.0.0
 * @createTime 2022/1/8 17:12
 */
public class BucketAlgorithm {

    private static final Logger logger = LoggerFactory.getLogger(BucketAlgorithm.class);

    /**
     * 流出的速度 每秒
     */
    private final int rate;

    /**
     * 桶大小
     */
    private final int bucketSize;

    /**
     * 刷新时间
     */
    private long refreshTime;

    /**
     * 当前桶容量大小
     */
    private final AtomicInteger currentSize = new AtomicInteger();

    public BucketAlgorithm(int rate, int bucketSize) {
        this.rate = rate;
        this.bucketSize = bucketSize;
    }

    /**
     * 刷新当前桶容量
     */
    private void refreshCurrentSize() {
        long now = System.currentTimeMillis();
        // 当前桶的容量等于 当前容量-（当前时间-上次记录时间）*速率
        currentSize.set((int) Math.max(0, currentSize.get() - (now - refreshTime) * rate));
        // 记录最近更新的时间
        refreshTime = now;
    }

    public synchronized boolean tryAcquire() {
        logger.info("当前线程信息：{} -------- 当前桶容量为 {}", Thread.currentThread().getId(), currentSize);
        refreshCurrentSize();
        if (currentSize.get() < bucketSize) {
            currentSize.set(currentSize.get() + 1);
            return true;
        } else {
            return false;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        BucketAlgorithm bucketAlgorithm = new BucketAlgorithm(1, 100);
        ThreadPoolExecutor pool = new ThreadPoolExecutor(5, 20, 10000, TimeUnit.SECONDS, new ArrayBlockingQueue<>(1000));
        for (int i = 0; i < 1000; i++) {
            pool.submit(new Thread(() -> logger.info(String.valueOf(bucketAlgorithm.tryAcquire()))));
        }
    }
}
